AWS CDKでBedrockを実行するLambdaを作ってみた

はじめに

今回は、AWS CDKでBedrockを実行するLambdaと実行の際に必要なリソースをまとめて作るサンプルコードを作ったのでご紹介します。CDK要素薄いですが、現状Bedrockで推論を動かすまでに最低限必要な要素を確認していただければと思います。

サンプルコード

事前準備

以下のブログを参考に利用するモデルを有効化してください。

実行環境

項目名	バージョン
AWS CDK	2.99.1
Finch	0.7.0
Bedrockのモデル	ai21.j2-ultra-v1

コードの紹介

現状最低限必要な項目に絞っています。具体的には最新のboto3を使うためのLambda Layerの設定、IAM Policy/Roleの実装、Lambdaの実装を行います。権限ついて試したところサンプルのモデルを使うにはbedrock:InvokeModelアクションのみで十分なようでした。細かい権限管理が行き届いていて良さそうです。

export class AwsCdkBedrockSampleStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // MEMO: Setup to use the latest version of boto3,
    //      which will not be necessary once Lambda's boto3 is up-to-date.
    const layerVersion = new LayerVersion(this, "Boto3LayerVersion", {
      code: Code.fromAsset(path.join(__dirname, "../src/layer/boto3.zip")),
      compatibleRuntimes: [Runtime.PYTHON_3_11],
    });

    const bedrockAccessPolicy = new PolicyStatement({
      effect: Effect.ALLOW,
      // See: https://docs.aws.amazon.com/ja_jp/service-authorization/latest/reference/list_amazonbedrock.html
      actions: ["bedrock:InvokeModel"],
      resources: ["*"],
    });

    const bedrockAccessRole = new Role(this, "BedrockAccessRole", {
      assumedBy: new ServicePrincipal("lambda.amazonaws.com"),
    });
    bedrockAccessRole.addToPolicy(bedrockAccessPolicy);

    new PythonFunction(this, "BedrockFunction", {
      entry: path.join(__dirname, "../src"),
      runtime: Runtime.PYTHON_3_11,
      index: "handler.py",
      handler: "handler",
      layers: [layerVersion],
      role: bedrockAccessRole,
      timeout: cdk.Duration.seconds(30),
    });
  }
}

最後に推論を呼び出すLambdaを以下で実装します。余談ですが公式サンプルのレスポンスがサンプルコードと異なっていたので、モデルによってレスポンスが若干異なるのかもです。

import boto3
import json
bedrock = boto3.client(service_name='bedrock-runtime',region_name='us-east-1')

modelId = 'ai21.j2-ultra-v1'
accept = 'application/json'
contentType = 'application/json'

def handler(event, context):
    request_message = event.get('request_message')
    body = json.dumps({
        "prompt": request_message,
        "maxTokens": 300,
        "temperature": 0.1,
        "topP": 0.9,
    })
    response = bedrock.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType)
    # HAK: Response body may vary by model type
    #      If the output is empty, check with the following code
    #      print(json.dumps(response_body))
    response_body = json.loads(response.get('body').read())
    result = response_body.get('completions')[0].get('data').get('text')
    response = {'result': result}
    return response

実行方法

# 今回はFinchを使うようにCDKを設定
% export CDK_DOCKER=finch

# 最新のBoto3を使うためのレイヤーを準備（Lambdaが最新のBoto3を使うようになれば不要）
% finch vm start
% finch run --entrypoint "" -v "$PWD":/var/task "public.ecr.aws/lambda/python:3.11.2023.09.27.10" /bin/sh -c "mkdir -p /tmp/python && pip3 install boto3 -t /tmp/python && cd /tmp && yum install -y zip && zip -r /var/task/boto3.zip ."
% mkdir src/layer
% mv boto3.zip src/layer/boto3.zip

# ローカル設定とデプロイ
% npm i
% npm run build
% npx cdk synth
% npx cdk deploy

# Lambda部分のみ再度デプロイする場合
% npx cdk deploy --hotswap

最新はReadme.mdで管理します。

動作確認

実際にGUI上でイベントを投げてLambdaを実行して確認してみます。

{
  "request_message": "What is CDK?"
}

すると以下のようにレスポンスが返ります。

所感

AWSだけで簡単に生成AIのAPI実行ができるようになったので少し試してみました。本題と関係ないですが、CDKのPythonFunction初めて使ってhotswapだとかなり早いので便利だなと感じました。今回のコードはBedrockを参考にしながら作成しました。サンプル程度の実装ですが参考になれば幸いです。今後需要あればBedrockの実行記録など周辺実装も追加したいと思います。

CX事業本部アーキテクトチームの佐藤智樹でした。

参考資料